//
// Created by ningjian on 2021/9/18.
//

#include "hscript_engine.hpp"
#include <regex>
#include <boost/algorithm/string.hpp>
#include <boost/algorithm/hex.hpp>
#include <algorithm>
#include <fstream>
#include <boost/lexical_cast.hpp>
#include "cmodel.h"

HScriptEngine::HScriptEngine()
{
    // set default value
    s_variable["p_bitcount"] = "0x400";
    s_variable["e_bitcount"] = "0x400";

    s_eval_fun = {
            {"set", std::bind(&HScriptEngine::eval_set, this, std::placeholders::_1)},
            {"saveblock", std::bind(&HScriptEngine::eval_save_block, this, std::placeholders::_1)},
            {"!*", std::bind(&HScriptEngine::eval_mod_mul, this, std::placeholders::_1)},
            {"!^", std::bind(&HScriptEngine::eval_mod_exp, this, std::placeholders::_1)},
            {"@*", std::bind(&HScriptEngine::eval_sum_mod_mul, this, std::placeholders::_1)},
            {"@^", std::bind(&HScriptEngine::eval_mul_mod_exp, this, std::placeholders::_1)},
    };
}
int HScriptEngine::Eval(std::string script_text)
{
    s_expression.clear();
    if(!parse_expression_to_vector(script_text)) {
        return -1;
    }

    for(auto &it : s_expression) {
        if(!eval_expression(it)) {
            return -1;
        }
    }

    return 0;
}

// 解析输入 拆分单个s表达式
bool HScriptEngine::parse_expression_to_vector(std::string& script_text)
{
    std::regex s_re("(\\([^)]+\\))");
    std::smatch matchResult;

    std::string m_script_text = script_text;
    boost::algorithm::replace_all(m_script_text, "\r", " ");
    boost::algorithm::replace_all(m_script_text, "\n", " ");

    string::const_iterator iterStart = m_script_text.begin();
    string::const_iterator iterEnd = m_script_text.end();
    while (std::regex_search(iterStart, iterEnd, matchResult, s_re))
    {
        s_expression.push_back(matchResult[0]);
        iterStart = matchResult[0].second;	//更新搜索起始位置,搜索剩下的字符串
    }

    return true;
}
// 解析单条表达式，并执行
bool HScriptEngine::eval_expression(std::string & s_text)
{
    boost::algorithm::replace_all(s_text, "\t", " ");
    boost::algorithm::replace_all(s_text, "  ", " ");
    boost::algorithm::replace_all(s_text, "  ", " ");

    std::string cmd = s_text;
    std::cout << "run:" << s_text << std::endl;

    boost::algorithm::replace_all(s_text, "(", "");
    boost::algorithm::replace_all(s_text, ")", "");

    std::vector<std::string> args;
    boost::algorithm::split(args,s_text,boost::is_any_of(" "));
    std::vector<std::string> new_args;
    for(auto &it : args){
        std::string ss =it;
        boost::algorithm::replace_all(ss, " ", "");
        if(!ss.empty()){
            new_args.push_back(it);
            //std::cout << it << std::endl;
        }
    }

    if( s_eval_fun.find(new_args.front()) == s_eval_fun.end()){
        return true;
    }

    if(!s_eval_fun[new_args.front()](new_args)) {
        std::cout << "eval expression failed!:" << cmd << std::endl;
        return false;
    }
    return true;

}

// set 命令
bool HScriptEngine::eval_set(std::vector<std::string> &args)
{
//    std::cout << "eval_set " << std::endl;

    if(args.size() < 3) {
        return false;
    }

    s_variable[args[1]] = args[2];

    return true;
}

bool HScriptEngine::eval_save_block(std::vector<std::string> &args)
{
    if(args.size() < 2) {
        return false;
    }
    if(s_blocks.size()  == 1){
        auto block = s_blocks.begin()->get_raw_data();
        save_vec_data(args[1]+".bin",block);
    }else {
        for(int i =0; i< s_blocks.size(); ++i) {
            auto block = s_blocks[i].get_raw_data();
            save_vec_data(args[1]+"_p"+std::to_string(i)+".bin",block);
        }
    }

    // save result data

    HomoBlock result_block;
    result_block.init(last_cal_mode);
    for(auto &it : s_frames) {
        result_block.add_frame(it);
    }
    auto result_data = result_block.get_raw_data(true);
    auto result_block_data = convert_result_number(result_data);
    save_vec_data(args[1]+"_result.bin",result_block_data);

    int p_bitcount = boost::lexical_cast<HexTo<int>>(s_variable["p_bitcount"].c_str());
    std::ofstream outFile(args[1]+"_result.txt", ios::out | ios::binary);
    std::string hex_a;
    boost::algorithm::hex(result_block_data.begin(), result_block_data.end(),std::back_inserter(hex_a));
    for(int i = 0; i < hex_a.size(); i+=4) {
        if( (i != 0)  && (i % (32)) == 0) {
            outFile << "\r\n";
        }
        outFile << hex_a.substr(i,4) << " ";
    }
    outFile << std::endl;
    outFile.close();

//    for(auto &it : s_result) {
//        std::string hex_a;
//        boost::algorithm::hex(it.get(), it.get()+p_bitcount/8, std::back_inserter(hex_a));
//        for(int i = 0; i < hex_a.size(); i+=4) {
//            if( (i != 0)  && (i % (32)) == 0) {
//                outFile << "\r\n";
//            }
//            outFile << hex_a.substr(i,4) << " ";
//        }
//        outFile << std::endl;
//    }
//    outFile.close();

    return true;
}

bool HScriptEngine::eval_mod_mul(std::vector<std::string> &args)
{
    return eval_mod_mul_and_exp(args, OP_CODE::OP_MUL);
}

bool HScriptEngine::eval_mod_exp(std::vector<std::string> &args)
{
    return eval_mod_mul_and_exp(args, OP_CODE::OP_EXP);

}

bool HScriptEngine::eval_mod_mul_and_exp(std::vector<std::string> &args, OP_CODE op)
{
    s_result.clear();
    s_blocks.clear();
    s_frames.clear();
    if((args.size()) < 4 || (args.size() % 2 != 0)) {
        std::cout << "args error!" <<std::endl;
        return false;
    }

    int p_bitcount = boost::lexical_cast<HexTo<int>>(s_variable["p_bitcount"].c_str());
    int e_bitcount = boost::lexical_cast<HexTo<int>>(s_variable["e_bitcount"].c_str());

    auto p = hex_2_array(args[1],p_bitcount / 8);

    // gen n r
    boost::shared_array<uint8_t> r(new uint8_t[p_bitcount / 8]);
    boost::shared_array<uint8_t> n(new uint8_t[128 / 8]);
    gen_n_prime_and_r_square((char *)p.get(),p_bitcount,(char *)n.get(),(char *)r.get());

    std::vector<std::pair<boost::shared_array<uint8_t>, boost::shared_array<uint8_t>>> numbers;
    for(int i = 2; i < args.size() ; i+=2) {
        auto a =  hex_2_array(args[i],p_bitcount / 8);
        if(op == OP_MUL_SUM || op == OP_MUL ) {
            auto b =  hex_2_array(args[i + 1],p_bitcount / 8);
            numbers.push_back(std::make_pair(a,b));
        }else {
            auto e =  hex_2_array(args[i + 1],e_bitcount / 8);
            numbers.push_back(std::make_pair(a,e));
        }

    }

    // 需要的pe个数
    s_blocks.clear();

    int now_calc_frame_index = 0;
    for(auto &it : numbers) {
        if(now_calc_frame_index % 80 == 0) {
            s_blocks.push_back(HomoBlock());
            if(op == OP_EXP || op == OP_MUL ) {
                last_cal_mode = CAL_MODE::CAL_MUL_OR_EXP;
                s_blocks.back().init(CAL_MODE::CAL_MUL_OR_EXP);
            }else if (op == OP_MUL_SUM) {
                last_cal_mode = CAL_MODE::CAL_MUL_SUM_ACCA;
                s_blocks.back().init(CAL_MODE::CAL_MUL_SUM_ACCA);
            }else if (op == OP_EXP_MUL) {
                last_cal_mode = CAL_MODE::CAL_EXP_MUL_ACCM;
                s_blocks.back().init(CAL_MODE::CAL_EXP_MUL_ACCM);
            }
            HomoFrame frame_setp;
            vector<uint8_t> number_p = array_2_vector(p, p_bitcount / 8);
            vector<uint8_t> number_r = array_2_vector(r, p_bitcount / 8);
            vector<uint8_t> number_n = array_2_vector(n, 128 / 8);
            //number_r.insert(number_r.end(),number_n.begin(),number_n.end());
            number_n.insert(number_n.end(),number_r.begin(),number_r.end());
            frame_setp.init(0x00,OP_CODE::OP_SET,get_pbc(p_bitcount),get_ebc(e_bitcount),number_p, number_n);
            s_blocks.back().add_frame(frame_setp);
        }

        HomoFrame frame_op;
        HomoFrame frame_result;
        vector<uint8_t> number_1 = array_2_vector(it.first, p_bitcount / 8);
        vector<uint8_t> number_2;
        if(op == OP_MUL_SUM || op == OP_MUL ) {
            number_2 = array_2_vector(it.second, p_bitcount / 8);
        }else {
            number_2 = array_2_vector(it.second, e_bitcount / 8);
        }

        frame_op.init(get_frame_id(now_calc_frame_index),op,get_pbc(p_bitcount),get_ebc(e_bitcount),number_1, number_2);

        s_blocks.back().add_frame(frame_op);

        if(s_variable["show"] == "true") {
            std::string hex_a;
            std::string hex_b;
            std::string hex_p;
            std::string hex_r;
            std::string hex_n;
            boost::algorithm::hex(it.first.get(), it.first.get()+p_bitcount/8, std::back_inserter(hex_a));
            if (OP_EXP == op || OP_EXP_MUL == op){
                boost::algorithm::hex(it.second.get(), it.second.get()+e_bitcount/8, std::back_inserter(hex_b));
            }else {
                boost::algorithm::hex(it.second.get(), it.second.get()+p_bitcount/8, std::back_inserter(hex_b));
            }

            boost::algorithm::hex(p.get(), p.get()+p_bitcount/8, std::back_inserter(hex_p));
            boost::algorithm::hex(r.get(), r.get()+p_bitcount/8, std::back_inserter(hex_r));
            boost::algorithm::hex(n.get(), n.get()+128/8, std::back_inserter(hex_n));

            std::cout << "A:" << hex_a << std::endl;
            std::cout << "B/E:" << hex_b << std::endl;
            std::cout << "P:" << hex_p << std::endl;
            std::cout << "R:" << hex_r << std::endl;
            std::cout << "N:" << hex_n << std::endl;
        }
        // calc result
        std::vector<uint8_t> result_number;
        if(OP_MUL== op || OP_MUL_SUM == op) {
            mod_mul_ssl(1, (char *)(it.first.get()),(char *)(it.second.get()), (char *)(p.get()),p_bitcount, (char *)(n.get()),(char *)(r.get()));
            s_result.push_back(it.first);
            result_number = convert_shared_array(it.first, p_bitcount/8);
        }else if (OP_EXP == op || OP_EXP_MUL == op){
            boost::shared_array<uint8_t> out(new uint8_t[p_bitcount/8]);
            mod_exp_ssl(1, (char *)(it.first.get()),(char *)(it.second.get()),e_bitcount, (char *)(p.get()),p_bitcount, (char *)(n.get()),(char *)(r.get()), (char *)out.get());
            s_result.push_back(out);
            result_number = convert_shared_array(out, p_bitcount/8);
        }

        std::vector<uint8_t> empty_number;
        if(op == OP_EXP || op == OP_MUL ) {
            frame_result.init(get_frame_id(now_calc_frame_index),op,get_pbc(p_bitcount),get_ebc(e_bitcount),result_number,empty_number);

        } else {
            frame_result.init(get_frame_id(now_calc_frame_index),op,get_pbc(p_bitcount),get_ebc(e_bitcount),empty_number,empty_number);
        }
        s_frames.push_back(frame_result);

        std::string hex_result;
        boost::algorithm::hex(s_result.back().get(), s_result.back().get()+p_bitcount/8, std::back_inserter(hex_result));
        std::cout << "GROUP:" << (get_frame_id(now_calc_frame_index) >> 4) << " PE:" << (get_frame_id(now_calc_frame_index) & 0x0f) <<"\nRESULT:"<< hex_result << std::endl;

        if(++now_calc_frame_index == 80) {
            now_calc_frame_index = 0;
        }
    }

    if(op == OP_MUL_SUM) {
        HomoFrame frame_end;
        vector<uint8_t> number_1;
        vector<uint8_t> number_2;
        frame_end.init(00, OP_MUL_SUM_ACCA,get_pbc(p_bitcount),get_ebc(e_bitcount),number_1,number_2);
        s_blocks.back().add_frame(frame_end);

        vector<uint8_t > rt = calc_result_sum(p);
        HomoFrame frame_rt;
        frame_rt.init(00, OP_MUL_SUM_ACCA,get_pbc(p_bitcount),get_ebc(e_bitcount),rt,number_2);
        s_frames.push_back(frame_rt);

    }else if(op == OP_EXP_MUL) {
        HomoFrame frame_end;
        vector<uint8_t> number_1;
        vector<uint8_t> number_2;
        frame_end.init(00, OP_EXP_MUL_ACCM,get_pbc(p_bitcount),get_ebc(e_bitcount),number_1,number_2);
        s_blocks.back().add_frame(frame_end);

        vector<uint8_t > rt = calc_result_mul(p);
        HomoFrame frame_rt;
        frame_rt.init(00, OP_EXP_MUL_ACCM,get_pbc(p_bitcount),get_ebc(e_bitcount),rt,number_2);
        s_frames.push_back(frame_rt);
    }

    return true;
}

std::vector<uint8_t> HScriptEngine::convert_result_number(std::vector<uint8_t> data){
    std::vector<uint8_t> result;
    for(int i = 0; i < data.size(); i+= 4) {
        result.push_back(data[i+3]);
        result.push_back(data[i+2]);
        result.push_back(data[i+1]);
        result.push_back(data[i+0]);
    }

    return result;
}

std::vector<uint8_t> HScriptEngine::convert_shared_array(boost::shared_array<uint8_t> number, int length)
{
    std::vector<uint8_t> result;
    for(int i = 0; i < length; ++i) {
        result.push_back(number[i]);
    }
    return result;
}


vector<uint8_t > HScriptEngine::calc_result_sum(boost::shared_array<uint8_t> p)
{
    vector<uint8_t > v_result;
    if(s_result.size() < 2) return v_result;
    int p_bitcount = boost::lexical_cast<HexTo<int>>(s_variable["p_bitcount"].c_str());
    BIGNUM *bn_a = BN_new();
    BIGNUM *bn_b = BN_new();
    BIGNUM *bn_r = BN_new();
    BN_CTX * ctx = BN_CTX_new();
    BN_bin2bn((const unsigned char *)s_result[0].get(),p_bitcount/8, bn_a);

    for(int i = 1; i < s_result.size(); ++i){
        BN_bin2bn((const unsigned char *)s_result[i].get(),p_bitcount/8, bn_b);
        BN_add(bn_r, bn_a, bn_b);
        BN_copy(bn_a,bn_r);
    }
    BN_bin2bn((const unsigned char *)p.get(),p_bitcount/8, bn_b);
    BN_mod(bn_r, bn_a, bn_b, ctx);

    boost::shared_array<uint8_t> result(new uint8_t[p_bitcount/8]);
    BN_bn2binpad(bn_r, (unsigned char *)result.get(),p_bitcount/8);

    BN_CTX_free(ctx);
    BN_free(bn_a);
    BN_free(bn_b);
    BN_free(bn_r);

    s_result.clear();
    s_result.push_back(result);

    v_result = array_2_vector(result, p_bitcount/8);
    array_2_vector(result, p_bitcount/8);

    std::string result_sum;
    boost::algorithm::hex(result.get(), result.get()+p_bitcount/8, std::back_inserter(result_sum));
    std::cout << "RESULT SUM:" << result_sum << std::endl;
    return v_result;
}

vector<uint8_t > HScriptEngine::calc_result_mul(boost::shared_array<uint8_t> p)
{
    vector<uint8_t > v_result;
    if(s_result.size() < 2) return v_result;
    int p_bitcount = boost::lexical_cast<HexTo<int>>(s_variable["p_bitcount"].c_str());
    BIGNUM *bn_a = BN_new();
    BIGNUM *bn_b = BN_new();
    BIGNUM *bn_r = BN_new();
    BN_CTX * ctx = BN_CTX_new();
    BN_bin2bn((const unsigned char *)s_result[0].get(),p_bitcount/8, bn_a);

    for(int i = 1; i < s_result.size(); ++i){
        BN_bin2bn((const unsigned char *)s_result[i].get(),p_bitcount/8, bn_b);
        BN_mul(bn_r, bn_a, bn_b,ctx);
        BN_copy(bn_a,bn_r);
    }
    BN_bin2bn((const unsigned char *)p.get(),p_bitcount/8, bn_b);
    BN_mod(bn_r, bn_a, bn_b, ctx);

    boost::shared_array<uint8_t> result(new uint8_t[p_bitcount/8]);
    BN_bn2binpad(bn_r, (unsigned char *)result.get(),p_bitcount/8);

    BN_CTX_free(ctx);
    BN_free(bn_a);
    BN_free(bn_b);
    BN_free(bn_r);

    s_result.clear();
    s_result.push_back(result);

    v_result = array_2_vector(result, p_bitcount/8);

    std::string result_mul;
    boost::algorithm::hex(result.get(), result.get()+p_bitcount/8, std::back_inserter(result_mul));
    std::cout << "RESULT MUL:" << result_mul << std::endl;

    return v_result;
}

bool HScriptEngine::eval_sum_mod_mul(std::vector<std::string> &args)
{
   return eval_mod_mul_and_exp(args, OP_MUL_SUM);
}

bool HScriptEngine::eval_mul_mod_exp(std::vector<std::string> &args)
{
    return eval_mod_mul_and_exp(args, OP_EXP_MUL);
}


void HScriptEngine::save_vec_data(std::string filename, vector<uint8_t> &data)
{
    std::ofstream outFile(filename, ios::out | ios::binary);
    for(auto &it : data) {
        outFile.write((char*)&it, 1);
    }
    outFile.close();
}

boost::shared_array<uint8_t>  HScriptEngine::hex_2_array(std::string hex, int array_size)
{
    boost::shared_array<uint8_t> result(new uint8_t[array_size]);
    for( int i =0; i< array_size; ++i){
        result[i] = 0;
    }

    boost::algorithm::replace_first(hex, "0x", "");
    if(hex.size() % 2 == 1)hex.insert(hex.begin(),'0');

    std::vector<uint8_t > bytes;
    boost::algorithm::unhex(hex, std::back_inserter(bytes));

    for(int i = array_size - bytes.size(); i < array_size; ++i) {
        result[i] = bytes[i - (array_size - bytes.size())];
    }

    return result;

}

vector<uint8_t > HScriptEngine::array_2_vector(boost::shared_array<uint8_t> arr, int arr_len)
{
    vector<uint8_t > result;
    for(int i =0 ; i< arr_len; ++i ){
        result.push_back(arr.get()[i]);
    }
    return result;
}
PBitCount HScriptEngine::get_pbc(int bitcount)
{
    switch (bitcount) {
        case 512: return PB_512;
        case 1024: return PB_1024;
        case 2048: return PB_2048;
        case 3072: return PB_3072;
        default: return PB_4096;
    }
}
EBitCount HScriptEngine::get_ebc(int bitcount)
{
    switch (bitcount) {
        case 64: return EB_64;
        case 128: return EB_128;
        case 256: return EB_256;
        case 512: return EB_512;
        case 1024: return EB_1024;
        case 2048: return EB_2048;
        case 3072: return EB_3072;
        default: return EB_4096;
    }
}

uint8_t  HScriptEngine::get_frame_id(int frame_index)
{
    frame_index = frame_index % 80 ;

    int group = frame_index / 10;
    int pe_index = frame_index % 10;

    return (group << 4) + pe_index;
}
